//
// Created by Tunm-Air13 on 2024/3/14.
//

#include "semantic_cls.h"
#include "log.h"
#include "serror.h"
#include "fixed_dict.h"
#include "labels.h"

namespace solex {

Indexes SemanticCls::getWordsIds(Words &words) {
    std::vector<int32_t> ids;
    ids.reserve(words.size());
    for (auto &word: words) {
        ids.push_back(m_interpreter_->getWordId(word));
    }

    return ids;
}

Words SemanticCls::cut(const Sentence &sentence) {
    std::vector<std::string> words;
    m_jieba_->Cut(sentence, words, m_use_hmm_);

    return words;
}

SemanticCls::SemanticCls() = default;

int32_t SemanticCls::Initialization(const std::string &folderPath) {
    std::string modelPath = folderPath + "/" + MODEL_PATH;
    std::string dictPath = folderPath + "/" + DICT_PATH;
    std::string hmmPath = folderPath + "/" + HMM_PATH;
    std::string userDictPath = folderPath + "/" + USER_DICT_PATH;
    std::string idfPath = folderPath + "/" + IDF_PATH;
    std::string stopWordPath = folderPath + "/" + STOP_WORD_PATH;

    m_interpreter_ = std::make_shared<fasttext::FastText>();
    try {
        m_interpreter_->loadModel(modelPath);
    } catch (const std::invalid_argument& e) {
        std::cerr << "Unable to open file: " << e.what() << std::endl;
        LOGE("Unable to open model file: %s", modelPath.c_str());

        return SERR_MODEL_FILE_ERROR;
    } catch (const std::exception& e) {
        std::cerr << "Load model error: " << e.what() << std::endl;
        LOGE("Load model error.");

        return SERR_MODEL_FILE_ERROR;
    }

    std::vector<std::string> needCheckList = {
            dictPath, hmmPath, userDictPath, idfPath, stopWordPath
    };

    // Check files
    for (const auto& filePath : needCheckList) {
        if (!checkFileExists(filePath)) {
            LOGE("The file does not exist: %s", filePath.c_str());
            return SERR_JIEBA_FILE_ERROR;
        }
    }

    // Load Jieba
    m_jieba_ = std::make_shared<cppjieba::Jieba>(
            dictPath,
            hmmPath,
            userDictPath,
            idfPath,
            stopWordPath
    );

    // Insert words
    for (auto &word: FIXED_DICT) {
        m_jieba_->InsertUserWord(word);
    }

    return SSUCCEED;
}

int32_t SemanticCls::Predict(const Sentence &sentence, SemanticCls::SemanticPredictResult &output) {
    auto words = cut(sentence);
    auto wordsIds = getWordsIds(words);
    fasttext::Predictions predictions;
    m_interpreter_->predict(1, wordsIds, predictions);
    auto idx = predictions[0].second;
    auto probability = std::exp(predictions[0].first);
    output.idx = idx;
    output.label = COMMAND_LABELS[idx].first;
    output.labelCh = COMMAND_LABELS[idx].second;
    output.cutResult = limonp::Join(words.begin(), words.end(), "/");
    output.probability = probability;
    output.cutWords = words;

    return SSUCCEED;
}


}   // namespace solex